package ltd.bugs.job.selector;

import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.ZooDefs;
import org.apache.zookeeper.ZooKeeper;
import org.apache.zookeeper.data.Stat;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;

import ltd.bugs.job.selector.cache.CacheKeys;
import ltd.bugs.job.selector.cache.LocalCache;

/**
 * @author zhengzhq
 * @Title: JobSelector
 * @Description: 任务选择器
 * @date 2018/11/23下午3:56
 */
public class JobSelector implements Watcher {

  private static Logger logger = LoggerFactory.getLogger(JobSelector.class);

  private LocalCache cache = new LocalCache();

  private SelectorConfig selectorConfig;
  private ZooKeeper zooKeeper = null;

  // 当前节点是否被选中
  private boolean currentNodeSelected = false;
  private Stat stat = new Stat();

  public JobSelector(SelectorConfig selectorConfig) {
    this.selectorConfig = selectorConfig;

    initZKClient();
  }

  private void initZKClient() {
    try {
      zooKeeper = new ZooKeeper(selectorConfig.getZkConnectString(), selectorConfig.getZkSessionTimeout(), this);

      // 增加钩子函数，确保在进程正常退出的情况下及时关闭zookeeper连接
      // 如果没有显式关闭，需要30秒后系统自动回收
      Runtime.getRuntime().addShutdownHook(new Thread(new Runnable() {
        @Override
        public void run() {
          if (zooKeeper != null) {
            try {
              zooKeeper.close();
            } catch (InterruptedException e) {
              logger.error("close zookeeper error", e);
            }
          }
        }
      }));
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
  }

  /**
   * 尝试获取job的运行权限
   */
  private boolean tryGetJobPermission() {
    try {
      // 判断节点是否存在
      Stat existStat = zooKeeper.exists(selectorConfig.getJobPath(), false);

      if (existStat != null) {// 节点存在，判断是否为当前节点
        byte[] data = zooKeeper.getData(selectorConfig.getJobPath(), false, stat);
        if (new String(data).equals(cache.getNodeValue())) {
          logger.info("current node is selected!");
          currentNodeSelected = true;
        } else {
          currentNodeSelected = false;
        }
      } else {// 节点不存在，尝试创建节点
        zooKeeper.create(selectorConfig.getJobPath(), cache.getNodeValue().getBytes(),
                ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL);

        logger.info("get job permission successfully!");

        // 当前节点创建成功，代表当前节点获得job的运行权限
        currentNodeSelected = true;
      }


      return currentNodeSelected;
    } catch (KeeperException e) {
      if (e.code().intValue() == KeeperException.Code.NODEEXISTS.intValue()) {
        logger.warn("get job permission fail, node already exists");
        currentNodeSelected = false;
      } else {
        logger.error("try get job permission error", e);
      }
    } catch (Exception e) {
      logger.error("try get job permission error", e);
    }

    return false;
  }

  public boolean isSelected() {
    // 缓存过期，需要重新尝试获取权限，并更新缓存
    if (cache.isExpired()) {
      currentNodeSelected = false;
      cache.set(CacheKeys.EXPIRE_TIME, System.currentTimeMillis() + selectorConfig.getCacheDuration());

      return tryGetJobPermission();
    }

    return currentNodeSelected;
  }

  @Override
  public void process(WatchedEvent watchedEvent) {
    switch (watchedEvent.getType()) {
      // 当前运行的节点下线，可以尝试创建node，正确job的运行权限
      case NodeDeleted:
        currentNodeSelected = false;
        tryGetJobPermission();
    }
  }
}
